package com.asr.service;

import com.asr.config.AsrProperties;
import com.asr.model.AsrConfig;
import com.asr.model.AsrTask;
import com.asr.model.ProcessProgress;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
@RequiredArgsConstructor
public class AsrTaskService {
    
    private final AsrProperties asrProperties;
    private final SimpMessagingTemplate messagingTemplate;
    private final ObjectMapper objectMapper;
    private final SshService sshService;
    
    // 存储任务状态
    private final Map<String, AsrTask> tasks = new ConcurrentHashMap<>();
    
    /**
     * 创建 ASR 任务
     */
    public AsrTask createTask(List<String> audioFiles, AsrConfig config, String outputDir) {
        String taskId = UUID.randomUUID().toString();
        
        AsrTask task = new AsrTask();
        task.setTaskId(taskId);
        task.setStatus("PENDING");
        task.setCreateTime(LocalDateTime.now());
        task.setConfig(config);
        task.setAudioFiles(audioFiles);
        task.setOutputDir(outputDir != null ? outputDir : asrProperties.getOutputDir());
        
        // 初始化文件状态
        for (String audioFile : audioFiles) {
            AsrTask.AudioFileStatus fileStatus = new AsrTask.AudioFileStatus();
            fileStatus.setFileName(Paths.get(audioFile).getFileName().toString());
            fileStatus.setStatus("PENDING");
            fileStatus.setProgress(0.0);
            fileStatus.setStage("pending");
            fileStatus.setMessage("等待处理");
            task.getFileStatuses().add(fileStatus);
        }
        
        tasks.put(taskId, task);
        log.info("创建任务: {}, 音频文件数: {}", taskId, audioFiles.size());
        
        return task;
    }
    
    /**
     * 获取任务信息
     */
    public AsrTask getTask(String taskId) {
        return tasks.get(taskId);
    }
    
    /**
     * 获取所有任务
     */
    public List<AsrTask> getAllTasks() {
        return new ArrayList<>(tasks.values());
    }
    
    /**
     * 异步执行 ASR 任务
     */
    @Async
    public void executeTask(String taskId) {
        AsrTask task = tasks.get(taskId);
        if (task == null) {
            log.error("任务不存在: {}", taskId);
            return;
        }
        
        task.setStatus("RUNNING");
        task.setStartTime(LocalDateTime.now());
        
        try {
            // 创建配置文件
            File configFile = createConfigFile(task);
            
            // 判断是本地执行还是远程执行
            if (asrProperties.getRemote() != null && asrProperties.getRemote().getEnabled()) {
                // 远程执行
                executeRemoteTask(task, configFile, taskId);
            } else {
                // 本地执行
                executeLocalTask(task, configFile, taskId);
            }
            
            // 清理配置文件
            configFile.delete();
            
        } catch (Exception e) {
            log.error("执行任务失败: {}", taskId, e);
            task.setStatus("FAILED");
            task.setErrorMessage(e.getMessage());
            task.setEndTime(LocalDateTime.now());
        }
        
        // 通知前端任务完成
        messagingTemplate.convertAndSend("/topic/task/" + taskId, task);
    }
    
    /**
     * 创建 Python 脚本配置文件
     */
    private File createConfigFile(AsrTask task) throws IOException {
        Map<String, Object> configData = new HashMap<>();
        
        // 任务信息
        configData.put("task_id", task.getTaskId());
        configData.put("audio_files", task.getAudioFiles());
        configData.put("output_dir", task.getOutputDir());
        
        // ASR 配置
        AsrConfig config = task.getConfig();
        Map<String, Object> asrConfig = new HashMap<>();
        asrConfig.put("speaker_segmentation_model", config.getSpeakerSegmentationModel());
        asrConfig.put("speaker_embedding_model", config.getSpeakerEmbeddingModel());
        asrConfig.put("asr_model_path", config.getAsrModelPath());
        asrConfig.put("asr_tokens_path", config.getAsrTokensPath());
        asrConfig.put("asr_model_type", config.getAsrModelType());
        
        if (config.getWhisperEncoder() != null) {
            asrConfig.put("whisper_encoder", config.getWhisperEncoder());
        }
        if (config.getWhisperDecoder() != null) {
            asrConfig.put("whisper_decoder", config.getWhisperDecoder());
        }
        
        asrConfig.put("num_speakers", config.getNumSpeakers());
        asrConfig.put("cluster_threshold", config.getClusterThreshold());
        asrConfig.put("language", config.getLanguage());
        asrConfig.put("num_processes", config.getNumProcesses());
        asrConfig.put("num_gpus", config.getNumGpus());
        asrConfig.put("num_threads", config.getNumThreads());
        asrConfig.put("sample_rate", config.getSampleRate());
        asrConfig.put("decoding_method", config.getDecodingMethod());
        asrConfig.put("use_itn", config.getUseItn());
        asrConfig.put("provider", config.getProvider());
        
        configData.put("config", asrConfig);
        
        // 写入临时文件
        File tempFile = File.createTempFile("asr_config_", ".json");
        objectMapper.writeValue(tempFile, configData);
        
        log.debug("创建配置文件: {}", tempFile.getAbsolutePath());
        return tempFile;
    }
    
    /**
     * 本地执行任务
     */
    private void executeLocalTask(AsrTask task, File configFile, String taskId) throws Exception {
        ProcessBuilder processBuilder = new ProcessBuilder(
            asrProperties.getPythonExecutable(),
            asrProperties.getScriptPath(),
            configFile.getAbsolutePath()
        );
        
        processBuilder.redirectErrorStream(true);
        Process process = processBuilder.start();
        
        // 读取输出
        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(process.getInputStream(), "UTF-8"))) {
            
            String line;
            while ((line = reader.readLine()) != null) {
                log.debug("Python 输出: {}", line);
                
                // 解析进度信息
                if (line.startsWith("PROGRESS:")) {
                    String progressJson = line.substring(9);
                    handleProgress(taskId, progressJson);
                }
                // 解析最终结果
                else if (line.startsWith("RESULT:")) {
                    String resultJson = line.substring(7);
                    handleResult(taskId, resultJson);
                }
            }
        }
        
        int exitCode = process.waitFor();
        
        if (exitCode == 0) {
            task.setStatus("COMPLETED");
            task.setEndTime(LocalDateTime.now());
            log.info("任务完成: {}", taskId);
        } else {
            task.setStatus("FAILED");
            task.setErrorMessage("Python 脚本执行失败，退出码: " + exitCode);
            log.error("任务失败: {}, 退出码: {}", taskId, exitCode);
        }
    }
    
    /**
     * 远程执行任务
     */
    private void executeRemoteTask(AsrTask task, File configFile, String taskId) throws Exception {
        // 上传配置文件到远程服务器
        String remoteConfigPath = "/tmp/asr_config_" + taskId + ".json";
        sshService.uploadFile(configFile.getAbsolutePath(), remoteConfigPath);
        
        // 构建远程执行命令
        String command;
        AsrProperties.Remote remote = asrProperties.getRemote();
        
        log.info("远程配置 - conda环境: {}", remote.getCondaEnv());

        // 如果配置了 Conda 环境名称 **并且** 提供了 Conda 绝对路径
        if (remote.getCondaEnv() != null && !remote.getCondaEnv().isEmpty() 
            && remote.getCondaPath() != null && !remote.getCondaPath().isEmpty()) {
            
            // **使用 conda run 方式执行命令 (推荐)**
            // 格式: /path/to/conda run -n <env> <python_exec> <script_path> <config_path>
            command = String.format(
                "%s run -n %s %s %s %s",
                remote.getCondaPath(), // /home/software/conda/anaconda3/bin/conda
                remote.getCondaEnv(),    // mama
                asrProperties.getPythonExecutable(),
                asrProperties.getScriptPath(),
                remoteConfigPath
            );
            
        } else if (remote.getCondaEnv() != null && !remote.getCondaEnv().isEmpty()) {
            // Fallback: 如果只配置了环境名，但没有配置绝对路径，仍然尝试原来的 `bash -lc`
            // 但这有风险，可能仍然失败。建议强制用户配置 conda-path。
            log.warn("已配置 Conda 环境 ({}) 但未配置 Conda 绝对路径，尝试使用 `bash -lc` 激活。", remote.getCondaEnv());
            command = String.format(
                "bash -lc 'conda activate %s && %s %s %s'",
                remote.getCondaEnv(),
                asrProperties.getPythonExecutable(),
                asrProperties.getScriptPath(),
                remoteConfigPath
            );
        } else {
            // 没有配置 conda，直接执行
            command = String.format("%s %s %s",
                asrProperties.getPythonExecutable(),
                asrProperties.getScriptPath(),
                remoteConfigPath
            );
        }
        
        log.info("执行远程命令: {}", command);
        
        // 执行远程命令并实时读取输出
        SshService.CommandResult result = sshService.executeCommandWithCallback(command, line -> {
            log.info("远程输出: {}", line);
            
            // 解析进度信息
            if (line.startsWith("PROGRESS:")) {
                String progressJson = line.substring(9);
                handleProgress(taskId, progressJson);
            }
            // 解析最终结果
            else if (line.startsWith("RESULT:")) {
                String resultJson = line.substring(7);
                handleResult(taskId, resultJson);
            }
        });
        
        // 清理远程配置文件
        try {
            sshService.executeCommand("rm -f " + remoteConfigPath);
        } catch (Exception e) {
            log.warn("清理远程配置文件失败: {}", e.getMessage());
        }
        
        log.info("远程命令执行完成，退出码: {}", result.getExitCode());
        log.info("标准输出: {}", result.getOutput());
        log.info("错误输出: {}", result.getError());
        
        if (result.isSuccess()) {
            task.setStatus("COMPLETED");
            task.setEndTime(LocalDateTime.now());
            log.info("远程任务完成: {}", taskId);
        } else {
            task.setStatus("FAILED");
            task.setErrorMessage("远程执行失败，退出码: " + result.getExitCode() + "\n错误: " + result.getError());
            log.error("远程任务失败: {}, 退出码: {}, 错误: {}", taskId, result.getExitCode(), result.getError());
        }
    }
    
    /**
     * 处理进度更新
     */
    private void handleProgress(String taskId, String progressJson) {
        try {
            ProcessProgress progress = objectMapper.readValue(progressJson, ProcessProgress.class);
            
            AsrTask task = tasks.get(taskId);
            if (task != null) {
                // 更新文件状态
                for (AsrTask.AudioFileStatus fileStatus : task.getFileStatuses()) {
                    if (fileStatus.getFileName().equals(progress.getFileName())) {
                        fileStatus.setStatus("PROCESSING");
                        fileStatus.setProgress(progress.getProgress());
                        fileStatus.setStage(progress.getStage());
                        fileStatus.setMessage(progress.getMessage());
                        fileStatus.setGpuId(progress.getGpuId());
                        
                        if ("completed".equals(progress.getStage())) {
                            fileStatus.setStatus("COMPLETED");
                        } else if ("error".equals(progress.getStage())) {
                            fileStatus.setStatus("FAILED");
                        }
                        break;
                    }
                }
            }
            
            // 通过 WebSocket 推送进度
            messagingTemplate.convertAndSend("/topic/progress/" + taskId, progress);
            
        } catch (Exception e) {
            log.error("解析进度信息失败: {}", progressJson, e);
        }
    }
    
    /**
     * 处理最终结果
     */
    private void handleResult(String taskId, String resultJson) {
        try {
            @SuppressWarnings("unchecked")
            Map<String, Object> result = objectMapper.readValue(resultJson, Map.class);
            
            AsrTask task = tasks.get(taskId);
            if (task != null) {
                Boolean success = (Boolean) result.get("success");
                
                if (Boolean.TRUE.equals(success)) {
                    task.setTotalDuration((Double) result.get("total_duration"));
                    task.setProcessingTime((Double) result.get("processing_time"));
                    task.setOverallRtf((Double) result.get("overall_rtf"));
                    
                    @SuppressWarnings("unchecked")
                    List<String> outputFiles = (List<String>) result.get("output_files");
                    if (outputFiles != null) {
                        task.setResultFiles(outputFiles);
                    }
                } else {
                    task.setErrorMessage((String) result.get("error"));
                }
            }
            
            log.info("任务 {} 结果: {}", taskId, result);
            
        } catch (Exception e) {
            log.error("解析结果信息失败: {}", resultJson, e);
        }
    }
    
    /**
     * 取消任务
     */
    public boolean cancelTask(String taskId) {
        AsrTask task = tasks.get(taskId);
        if (task != null && "RUNNING".equals(task.getStatus())) {
            task.setStatus("CANCELLED");
            task.setEndTime(LocalDateTime.now());
            log.info("取消任务: {}", taskId);
            return true;
        }
        return false;
    }
    
    /**
     * 删除任务
     */
    public boolean deleteTask(String taskId) {
        AsrTask task = tasks.remove(taskId);
        if (task != null) {
            log.info("删除任务: {}", taskId);
            return true;
        }
        return false;
    }
}
